import numpy as np
import pandas as pd

results = {
    "results-imagenet.csv": [
        "results-imagenet-real.csv",
        "results-imagenetv2-matched-frequency.csv",
        "results-sketch.csv",
    ],
    "results-imagenet-a-clean.csv": [
        "results-imagenet-a.csv",
    ],
    "results-imagenet-r-clean.csv": [
        "results-imagenet-r.csv",
    ],
}


def diff(base_df, test_csv):
    base_models = base_df["model"].values
    test_df = pd.read_csv(test_csv)
    test_models = test_df["model"].values

    rank_diff = np.zeros_like(test_models, dtype="object")
    top1_diff = np.zeros_like(test_models, dtype="object")
    top5_diff = np.zeros_like(test_models, dtype="object")

    for rank, model in enumerate(test_models):
        if model in base_models:
            base_rank = int(np.where(base_models == model)[0])
            top1_d = test_df["top1"][rank] - base_df["top1"][base_rank]
            top5_d = test_df["top5"][rank] - base_df["top5"][base_rank]

            # rank_diff
            if rank == base_rank:
                rank_diff[rank] = f"0"
            elif rank > base_rank:
                rank_diff[rank] = f"-{rank - base_rank}"
            else:
                rank_diff[rank] = f"+{base_rank - rank}"

            # top1_diff
            if top1_d >= 0.0:
                top1_diff[rank] = f"+{top1_d:.3f}"
            else:
                top1_diff[rank] = f"-{abs(top1_d):.3f}"

            # top5_diff
            if top5_d >= 0.0:
                top5_diff[rank] = f"+{top5_d:.3f}"
            else:
                top5_diff[rank] = f"-{abs(top5_d):.3f}"

        else:
            rank_diff[rank] = ""
            top1_diff[rank] = ""
            top5_diff[rank] = ""

    test_df["top1_diff"] = top1_diff
    test_df["top5_diff"] = top5_diff
    test_df["rank_diff"] = rank_diff

    test_df["param_count"] = test_df["param_count"].map("{:,.2f}".format)
    test_df.sort_values(["top1", "top5", "model"], ascending=[False, False, True], inplace=True)
    test_df.to_csv(test_csv, index=False, float_format="%.3f")


for base_results, test_results in results.items():
    base_df = pd.read_csv(base_results)
    base_df.sort_values(["top1", "top5", "model"], ascending=[False, False, True], inplace=True)
    for test_csv in test_results:
        diff(base_df, test_csv)
    base_df["param_count"] = base_df["param_count"].map("{:,.2f}".format)
    base_df.to_csv(base_results, index=False, float_format="%.3f")
